import asyncio
import os
from pydantic_ai import Agent, Tool


os.environ["DEEPSEEK_API_KEY"] = "sk-b7ad23e07ace4e96ab304ae01b45ae64"  # 仅用于演示，生产环境不应硬编码

# 定义工具函数（模拟API）
def get_sales_data(product_id: str, start_date: str, end_date: str) -> dict:
    return {"sales": 1500, "trend": "up 20%"}

def analyze_data(raw_data: dict) -> str:
    return f"报告：产品销量{raw_data['sales']}件，增长趋势{raw_data['trend']}"

def send_email(report: str, recipient: str) -> str:
    return f"邮件已发送至{recipient}：{report}"

# 创建Agent链
agent = Agent("deepseek:deepseek-chat",
               tools=[Tool(get_sales_data), Tool(analyze_data), Tool(send_email)],               
               system_prompt="执行任务时实时报告工具调用状态"
               )  # 工具直接传入

# 异步流式处理 - 兼容pydantic_ai 0.4.0
async def main():
    try:
        # 使用旧版本兼容的流式处理方式
        run_result = await agent.run(
            "分析产品P123在2025-07-01到2025-07-10的销售数据并发送报告至admin@example.com"
            # stream=True  # 显式启用流式输出
        )
        
        # 处理流式输出
        async for chunk in run_result.output:
            if chunk.type == "tool_invocation":
                print(f"\n工具调用：正在执行 {chunk.tool_name}()")
                print("-" * 40)
            elif chunk.type == "message_delta":
                print(chunk.content, end="", flush=True)
                
        # 打印最终结果
        if run_result.final_response:
            print("\n\n最终结果:", run_result.final_response)
            
    except Exception as e:
        print(f"执行过程中发生错误: {e}")

# 运行异步任务
asyncio.run(main())